package com.dark.aop.sleuth;

import org.slf4j.MDC;

import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.UUID;

/**
 * ClassName: TraceUtil
 * Description:
 *
 * @author chengluchao
 * @date: 2021/2/1 20:36
 * @since 1.1.12
 */
public abstract class TraceUtil {

    /**
     * 设置请求，响应头名字
     */
    private static final String REQUEST_ID_HEADER = "RequestId";
    /**
     * 设置到 MDC 里的key
     */
    public static final String TRACE_ID = "TraceId";

    public static final String SPAN_ID = "SpanId";


    /**
    /**
     * 刷新  TraceId and SpanId
     * <p>刷新数据</p>
     * <p>慎用：此方法会重置TraceId</p>
     * <p>使用场景</p>
     * <i>当请求不能进行spanId初始化时：</i>
     * <i1>MQ的回调</i1>
     * <i>定时任务的启动</i>
     */
    public static void refresh() {
        String spanId = generateId();
        TraceUtil.setSpanId(spanId);
        TraceUtil.setTraceId(spanId);
    }

    /**
     * 配置requestId
     *
     * @param request
     * @param response
     */
    protected static void configRequestId(ServletRequest request, ServletResponse response) {
        HttpServletRequest httpRequest = (HttpServletRequest) request;
        HttpServletResponse httpResponse = (HttpServletResponse) response;
        // span id in app
        String spanId = generateId();

        String traceId = httpRequest.getHeader(REQUEST_ID_HEADER);
        if (traceId == null || "".equals(traceId)) {
            traceId = spanId;
        }

        MDC.put(SPAN_ID, spanId);
        MDC.put(TRACE_ID, traceId);
        httpResponse.setHeader(REQUEST_ID_HEADER, traceId);
    }

    protected static String generateId() {
        return UUID.randomUUID().toString().replaceAll("-", "").substring(8, 24);
    }

    public static void setTraceId(String traceId) {
        MDC.put(TRACE_ID, traceId);
        if (getSpanId() == null) {
            setSpanId(generateId());
        }
    }

    static void setSpanId(String spanId) {
        MDC.put(SPAN_ID, spanId);
    }

    static void clear() {
        MDC.clear();
    }

    public static String getTraceId() {
        if (MDC.get(TRACE_ID) == null) {
            refresh();
        }

        return MDC.get(TRACE_ID);
    }

    public static String getSpanId() {
        if (MDC.get(SPAN_ID) == null) {
            setSpanId(generateId());
        }
        return MDC.get(SPAN_ID);
    }

    private TraceUtil() {
    }
}
